import type { ParseUniversalResult } from '../universalParser.ts';
import type { UniversalFlowNode, UniversalFlowEdge, SourceFormat, NodeSemanticType, UniversalFlowNodeData } from '@/models/universalFlow.ts';
import { v4 as uuidv4 } from 'uuid';
import { babelService } from '@/services/babelService';
import * as t from '@babel/types'; // Import all types under 't' namespace

interface ConversionContext {
    nodes: UniversalFlowNode[];
    edges: UniversalFlowEdge[];
    format: SourceFormat;
    level: number;
    path: string;
    order: number;
}

// Function to determine the semantic type based on Babel Node
function getSemanticType(node: t.Node | null): NodeSemanticType | null {
    if (!node) return null;

    if (t.isObjectExpression(node)) return 'object';
    if (t.isArrayExpression(node) || t.isTSArrayType(node)) return 'array';
    if (t.isStringLiteral(node) || t.isNumericLiteral(node) || t.isBooleanLiteral(node) || t.isNullLiteral(node) || t.isBigIntLiteral(node) || t.isDecimalLiteral(node)) return 'value';
    if (t.isIdentifier(node) && node.name === 'undefined') return 'value'; // Handle undefined identifier
    if (t.isIdentifier(node)) return 'variableReference';
    if (t.isCallExpression(node)) return 'functionCall';
    if (t.isMemberExpression(node)) return 'memberExpression';
    if (t.isBinaryExpression(node)) return 'binaryExpression';
    if (t.isUnaryExpression(node)) return 'unaryExpression';
    if (t.isConditionalExpression(node)) return 'conditionalExpression';
    if (t.isTemplateLiteral(node)) return 'templateLiteral';
    if (t.isNewExpression(node)) return 'newExpression';
    // Add more mappings as needed...

    // Fallback or ignore other types?
    console.warn(`[Babel Parser] Unhandled Babel node type for semantic mapping: ${node.type}`);
    return null; // Or a default like 'value' if appropriate?
}

// Recursive function to transform Babel AST
async function transformBabelNode(babelNode: t.Node | null | undefined, parentId: string, context: ConversionContext): Promise<string | null> {
    if (!babelNode || t.isSpreadElement(babelNode)) {
        return null;
    }

    const nodeId = uuidv4();
    const nodeSemanticType = getSemanticType(babelNode);

    if (nodeSemanticType === null) {
        return null;
    }

    const nodeData: Partial<UniversalFlowNodeData> = {
        name: '',
        format: context.format,
        level: context.level,
        path: context.path,
        order: context.order,
        nodeSemanticType: nodeSemanticType,
        metadata: {},
        attributes: {},
    };

    const childrenIds: string[] = [];

    // --- Process Node Content and Children --- 
    switch (babelNode.type) {
        case 'ObjectExpression':
            let propOrder = 0;
            for (const prop of babelNode.properties) {
                if (t.isObjectProperty(prop) && !prop.computed) {
                    let key: string | null = null;
                    if (t.isIdentifier(prop.key)) key = prop.key.name;
                    else if (t.isStringLiteral(prop.key)) key = prop.key.value;
                    else if (t.isNumericLiteral(prop.key)) key = String(prop.key.value);

                    if (key !== null) {
                        const childContext: ConversionContext = { ...context, level: context.level + 1, path: `${context.path ? context.path + '.' : ''}${key}`, order: propOrder++ };
                        const childId = await transformBabelNode(prop.value as t.Expression, nodeId, childContext);
                        if (childId) {
                            const childNode = context.nodes.find(n => n.id === childId);
                            if (childNode?.data) childNode.data.name = key;
                            childrenIds.push(childId);
                        }
                    }
                }
            }
            break;

        case 'ArrayExpression':
            let elementOrder = 0;
            for (const element of babelNode.elements) {
                if (element) {
                    const childContext: ConversionContext = { ...context, level: context.level + 1, path: `${context.path}[${elementOrder}]`, order: elementOrder++ };
                    const childId = await transformBabelNode(element, nodeId, childContext);
                    if (childId) {
                        childrenIds.push(childId);
                    }
                }
            }
            break;

        // --- Value Literals --- 
        case 'StringLiteral':
        case 'NumericLiteral':
        case 'BooleanLiteral':
            nodeData.value = babelNode.value;
            break;
        case 'NullLiteral':
            nodeData.value = null;
            break;
        case 'BigIntLiteral':
            nodeData.value = babelNode.value.toString() + 'n';
            break;
        case 'DecimalLiteral':
            nodeData.value = parseFloat(babelNode.value);
            break;
        case 'Identifier':
            if (babelNode.name === 'undefined') {
                nodeData.value = undefined;
            } else {
                nodeData.metadata!.variableName = babelNode.name;
            }
            break;

        // --- Other Expression Types (Only populate metadata if needed for display/info) --- 
        case 'CallExpression':
            if (t.isIdentifier(babelNode.callee)) {
                nodeData.metadata!.functionName = babelNode.callee.name;
            }
            break;
        case 'BinaryExpression':
        case 'UnaryExpression':
            nodeData.metadata!.operator = babelNode.operator;
            break;
        case 'NewExpression':
            if (t.isIdentifier(babelNode.callee)) {
                nodeData.metadata!.constructorName = babelNode.callee.name;
            }
            break;
        case 'TemplateLiteral':
        case 'MemberExpression':
        case 'ConditionalExpression':
            break;

        default:
            break;
    }

    // --- Create Universal Flow Node --- 
    const universalNode: UniversalFlowNode = {
        id: nodeId,
        type: 'universal',
        position: { x: 0, y: 0 },
        data: {
            sourceElement: babelNode,
            name: nodeData.name!,
            content: undefined,
            attributes: nodeData.attributes!,
            path: nodeData.path!,
            format: nodeData.format!,
            nodeSemanticType: nodeData.nodeSemanticType!,
            childrenIds: childrenIds.length > 0 ? childrenIds : undefined,
            value: nodeData.value,
            order: nodeData.order!,
            level: nodeData.level!,
            metadata: Object.keys(nodeData.metadata!).length > 0 ? nodeData.metadata : undefined
        }
    };
    context.nodes.push(universalNode);

    // --- Create Edge from Parent --- 
    context.edges.push({
        id: `e-${parentId}-${nodeId}`,
        source: parentId,
        target: nodeId,
        type: 'smoothstep',
        data: { relationshipType: 'parent-child' }
    });

    return nodeId;
}

/**
 * Parses JavaScript or TypeScript code using Babel and converts the AST.
 * Focuses on the content within `export default { ... }` or equivalent variable export.
 */
export async function parseBabelToUniversal(
    content: string,
    format: 'javascript' | 'typescript'
): Promise<ParseUniversalResult> {
    let ast: t.File | null = null;
    let success = false;
    let error: string | undefined = undefined;
    const nodes: UniversalFlowNode[] = [];
    const edges: UniversalFlowEdge[] = [];
    let rootId: string | null = null;

    try {
        console.log(`[Babel Parser] Attempting to parse ${format} content...`);
        ast = await babelService.parseCode(content);

        // --- ADDED: Robust AST structure checks ---
        if (!ast) {
            throw new Error('Babel parsing returned null AST.');
        }
        if (!ast.program) {
             throw new Error('Babel AST is missing the \'program\' property.');
        }
        if (!ast.program.body || !Array.isArray(ast.program.body)) {
            // Log the structure for debugging if body is missing or not an array
            console.error('[Babel Parser] Babel AST program body is missing or not an array:', ast.program);
            throw new Error('Babel AST program body is missing or not an array.');
        }
        // --- End Checks ---

        // --- Find the default export (Now safe to access body) ---
        let rootExportExpression: t.Expression | null = null;
        let exportNode: t.ExportDefaultDeclaration | undefined;

        // Find export default declaration directly
        exportNode = ast.program.body.find( // Accessing body is now safer
            (node): node is t.ExportDefaultDeclaration => t.isExportDefaultDeclaration(node)
        );

        if (exportNode) {
            const declaration = exportNode.declaration;
            if (t.isObjectExpression(declaration)) {
                // Case 1: export default { ... }
                rootExportExpression = declaration;
            } else if (t.isIdentifier(declaration)) {
                // Case 2: export default variable;
                const exportName = declaration.name;
                console.log(`[Babel Parser] Found export default identifier: ${exportName}`);
                // Find the variable declaration for this identifier
                for (const item of ast.program.body) {
                    if (t.isVariableDeclaration(item)) {
                        for (const declarator of item.declarations) {
                            if (t.isIdentifier(declarator.id) && declarator.id.name === exportName) {
                                if (declarator.init && t.isObjectExpression(declarator.init)) {
                                    rootExportExpression = declarator.init;
                                    console.log(`[Babel Parser] Found variable declaration for ${exportName} with ObjectExpression.`);
                                    break;
                                } else {
                                    error = `Variable '${exportName}' is exported but not initialized with an object literal.`;
                                    break;
                                }
                            }
                        }
                    }
                    if (rootExportExpression || error) break;
                }
                 if (!rootExportExpression && !error) {
                     error = `Could not find variable declaration for exported identifier '${exportName}'.`;
                }
            } else {
                // Handle other direct export types if needed (e.g., ArrayExpression)
                 if (t.isArrayExpression(declaration)) {
                     // If exporting an array directly is supported, assign it
                     rootExportExpression = declaration;
                     console.log(`[Babel Parser] Found direct export of an ArrayExpression.`);
                 } else {
                    error = `Unsupported default export type: ${declaration.type}. Only object/array literals or variables initialized with them are supported.`;
                 }
            }
        } else {
            error = 'Could not find default export (`export default ...`).';
        }

        // --- Process the found export expression --- 
        if (rootExportExpression) {
             // Create the root node for the exported object/array
            rootId = uuidv4();
            const rootNode: UniversalFlowNode = {
                id: rootId,
                type: 'universal',
                position: { x: 0, y: 0 },
                data: {
                    sourceElement: rootExportExpression,
                    name: '[root]',
                    format: format,
                    level: 0,
                    path: '',
                    order: 0,
                    nodeSemanticType: t.isObjectExpression(rootExportExpression) ? 'object' : 'array',
                    attributes: {},
                    metadata: {}
                }
            };
            nodes.push(rootNode);

            const context: ConversionContext = { nodes, edges, format, level: 0, path: '', order: 0 };
            const childrenIds: string[] = [];

            // Start transformation from the actual exported expression
            if (t.isObjectExpression(rootExportExpression)) {
                 let propOrder = 0;
                 for (const prop of rootExportExpression.properties) {
                     if (t.isObjectProperty(prop) && !prop.computed) {
                         let key: string | null = null;
                         if (t.isIdentifier(prop.key)) key = prop.key.name;
                         else if (t.isStringLiteral(prop.key)) key = prop.key.value;

                         if (key !== null) {
                             const childContext: ConversionContext = { ...context, level: 1, path: key, order: propOrder++ };
                             const childId = await transformBabelNode(prop.value as t.Expression, rootId, childContext);
                             if (childId) {
                                 const childNode = context.nodes.find(n => n.id === childId);
                                 if (childNode?.data) childNode.data.name = key;
                                 childrenIds.push(childId);
                             }
                         }
                     }
                 }
            } else if (t.isArrayExpression(rootExportExpression)) {
                 let elementOrder = 0;
                 for (const element of rootExportExpression.elements) {
                      if (element) {
                         const childContext: ConversionContext = { ...context, level: 1, path: `[${elementOrder}]`, order: elementOrder++ };
                         const childId = await transformBabelNode(element, rootId, childContext);
                         if (childId) {
                             childrenIds.push(childId);
                         }
                      }
                 }
            }

            // Assign children to the root node
            if (rootNode.data) {
                rootNode.data.childrenIds = childrenIds.length > 0 ? childrenIds : undefined;
            }
            success = true;

        } else {
            // Set error only if no other error has been set yet
            if (!error) { 
                error = `Failed to find or process a valid default export expression in the script.`;
                console.warn(`[Babel Parser] ${error}`);
                success = false; 
            }
        }

    } catch (parseOrProcessError: any) {
        console.error(`[Babel Parser] Error during parsing or processing:`, parseOrProcessError);
        error = parseOrProcessError.message || 'Babel parsing or processing failed.';
        success = false;
    }

    // Return result
    return {
        success,
        nodes,
        edges,
        rootId,
        rawDocument: ast,
        error
    };
} 